import os
from Crypto.Cipher import PKCS1_OAEP
from Crypto.PublicKey import RSA
import base64
import secrets

def load_rsa_key():
    base_dir = os.path.dirname(os.path.abspath(__file__)).replace('\\', '/') + "/../../"
    key_path = base_dir + "data/client/rsa/"
    sk1,sk2 = None,None
    with open(key_path + "sk1.pem", "rb") as key_file:
        sk1 = RSA.import_key(key_file.read())
    with open(key_path + "sk2.pem", "rb") as key_file:
        sk2 = RSA.import_key(key_file.read())
    return sk1,sk2

def work():
    base_dir = os.path.dirname(os.path.abspath(__file__)).replace('\\', '/') + "/../../data/client/"
    sk1,sk2 = load_rsa_key()
    cipher1 = PKCS1_OAEP.new(sk1)
    cipher2 = PKCS1_OAEP.new(sk2)
    with open(base_dir+"key_enc.txt", 'r') as fin, open(base_dir+"key_dec1.txt", 'w') as fout:
        for line in fin:
            try:
                dec_line = cipher1.decrypt(base64.b64decode(line.strip())).hex()
                fout.write(dec_line + "\n")
            except Exception as e:
                fout.write(secrets.token_hex(16)+ "\n")

    with open(base_dir+"key_enc.txt", 'r') as fin, open(base_dir+"key_dec2.txt", 'w') as fout:
        for line in fin:
            try:
                dec_line = cipher2.decrypt(base64.b64decode(line.strip())).hex()
                fout.write(dec_line + "\n")
            except Exception as e:
                fout.write(secrets.token_hex(16) + "\n")

    print("Decryption done!")
if __name__ == '__main__':
    work()

